# ===== 11-1.1 =====
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.tree import DecisionTreeRegressor
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score

# 设置中文字体
plt.rcParams["font.sans-serif"] = ["SimSong"]
plt.rcParams["axes.unicode_minus"] = False

# 原始数据加载和处理
gdp_df = pd.read_csv("./data/gdp_melted.csv")
bj_gdp = gdp_df[gdp_df["Region"] == "北京市"].copy()

# ===== 11-1.2 =====
# 准备特征和标签
X = bj_gdp[["Year"]]  # 使用Year和Population作为特征
y = bj_gdp["GDP"]  # 使用GDP作为标签

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)


# ===== 11-1.3 =====
# 初始化并训练随机森林模型
reg_model = RandomForestRegressor(random_state=42)
reg_model.fit(X_train, y_train)

# 预测
y_pred = reg_model.predict(X_test)

# ===== 11-1.4 =====
# 初始化并训练决策树模型
reg_model = DecisionTreeRegressor(random_state=42)
reg_model.fit(X_train, y_train)

# 预测
y_pred = reg_model.predict(X_test)

# ===== 11-1.5 =====
# 模型验证
mse = mean_squared_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)
print(f"MSE={mse:.2f}, R2={r2:.2f}")


print(f"呼和浩特人口: {get_data('呼和浩特市')}")